from utils.logger import setup_logger
from datasets import make_dataloader
from model import make_model
from solver import make_optimizer
from solver.scheduler_factory import create_scheduler
from loss import make_loss
from processor import do_train
import random
import torch
import numpy as np
import os
import argparse
from config import cfg
import time
import glob
os.environ["TF_ENABLE_ONEDNN_OPTS"] = "0"

#set_seed 函数通过设置随机种子（torch, numpy, random）确保实验可重复性，并启用 CuDNN 的确定性模式。
def set_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = True

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="ReID Baseline Training")
    parser.add_argument(
        "--config_file", default="", help="path to config file", type=str
    )
    parser.add_argument("opts", help="Modify config options using the command-line", default=None,
                        nargs=argparse.REMAINDER)
    parser.add_argument("--local_rank", default=0, type=int)
    args = parser.parse_args()

    if args.config_file != "":
        cfg.merge_from_file(args.config_file)
        print(f"Loaded configuration file: {args.config_file}")
        print(f"cfg.MODEL.NAME after merge: {cfg.MODEL.NAME}")
    cfg.merge_from_list(args.opts)
    print(f"cfg.MODEL.NAME after opts: {cfg.MODEL.NAME}")
    cfg.freeze()
#     if __name__ == '__main__':
#         parser = argparse.ArgumentParser(description="ReID Baseline Training")
#         parser.add_argument(
#             "--config_file", default="", help="path to config file", type=str
#         )
#         parser.add_argument("opts", help="Modify config options using the command-line", default=None,
#                             nargs=argparse.REMAINDER)
#         parser.add_argument("--local_rank", default=0, type=int)
#         args = parser.parse_args()
#
#         if args.config_file != "":
#             cfg.merge_from_file(args.config_file)
#             print(f"Loaded configuration file: {args.config_file}")
#             print(f"cfg.MODEL.NAME after merge: {cfg.MODEL.NAME}")
#         cfg.merge_from_list(args.opts)
#         print(f"cfg.MODEL.NAME after opts: {cfg.MODEL.NAME}")
#
#         # 自动生成递增的训练编号
#         base_output_dir = cfg.OUTPUT_DIR
#         existing_dirs = glob.glob(os.path.join(base_output_dir, "train*"))
#         train_id = len(existing_dirs) + 1
#         new_output_dir = os.path.join(base_output_dir, f"train{train_id}")
#         cfg.OUTPUT_DIR = new_output_dir
#         print(f"Updated OUTPUT_DIR: {cfg.OUTPUT_DIR}")
#
#         cfg.freeze()
#
#         # 确保新目录存在
#         if not os.path.exists(cfg.OUTPUT_DIR):
#             os.makedirs(cfg.OUTPUT_DIR)
#



    # CUDA 调试日志
    print(f"CUDA available: {torch.cuda.is_available()}, Device count: {torch.cuda.device_count()}, Device name: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'None'}")
    print(f"CUDA_VISIBLE_DEVICES: {os.environ.get('CUDA_VISIBLE_DEVICES', 'Not set')}")
    try:
        test_tensor = torch.randn(1, 10).cuda()
        print(f"CUDA test passed: Tensor on {test_tensor.device}")
    except Exception as e:
        print(f"CUDA test failed: {e}")

    device = torch.device('cuda' if cfg.MODEL.DEVICE == 'cuda' and torch.cuda.is_available() else 'cpu')
    print(f"Using device: {device}")
    if device.type == 'cpu':
        print("WARNING: Training on CPU, this will be slow. Check CUDA installation.")

    set_seed(cfg.SOLVER.SEED)

    if cfg.MODEL.DIST_TRAIN:
        torch.cuda.set_device(args.local_rank)
        torch.distributed.init_process_group(backend='nccl', init_method='env://')

    output_dir = cfg.OUTPUT_DIR
    if output_dir and not os.path.exists(output_dir):
        os.makedirs(output_dir)

    logger = setup_logger("transreid", output_dir, if_train=True)
    logger.info("Saving model in the path :{}".format(cfg.OUTPUT_DIR))
    logger.info(args)

    if args.config_file != "":
        logger.info("Loaded configuration file {}".format(args.config_file))
        with open(args.config_file, 'r', encoding='utf-8') as cf:
            config_str = "\n" + cf.read()
            logger.info(config_str)
    logger.info("Running with config:\n{}".format(cfg))

    train_loader, train_loader_normal, val_loader, num_query, num_classes, camera_num, view_num = make_dataloader(cfg)

    model = make_model(cfg, num_class=num_classes, camera_num=camera_num, view_num=view_num)
    model = model.to(device)

    loss_func, center_criterion = make_loss(cfg, num_classes=num_classes)

    optimizer, optimizer_center = make_optimizer(cfg, model, center_criterion)

    scheduler = create_scheduler(cfg, optimizer)

    do_train(
        cfg,
        model,
        center_criterion,
        train_loader,
        val_loader,
        optimizer,
        optimizer_center,
        scheduler,
        loss_func,
        num_query,
        args.local_rank,
        camera_num,  # <--- 添加这个
        view_num  # <--- 添加这个
    )
